import torch.utils.data as data
from torchvision import transforms 
import torch
import torch.nn.functional as F
import yaml
import Evaluation
import random
import h5py
#transforms.Normalize 会使得psnrin下降
path='./default.yaml'
with open(path, 'r') as f:
    cfg = yaml.load(f, Loader=yaml.Loader)
def cut_picture(M1,M2,M3,M4, dH=128,dW=128):    #batch==1   输入一个数组
    sizes=M1[0].size()
    outs1=[]
    outs2=[]
    outs3=[]
    outs4=[]
    H=sizes[-2]
    W=sizes[-1]

    x=random.randint(0,H-dH)
    y=random.randint(0,W-dW)
    for i1 in M1:
        outs1.append(i1[...,x:x+dH,y:y+dW])
    for i2 in M2:
        outs2.append(i2[...,x:x+dH,y:y+dW])
    for i3 in M3:
        outs3.append(i3[...,x:x+dH,y:y+dW])
    for i4 in M4:
        outs4.append(i4[...,x:x+dH,y:y+dW])
    return  outs1,outs2,outs3,outs4
class Database(data.Dataset):
    def __init__(self,flag):
         super( Database, self).__init__()
         if flag=='train':
            self.cames1= self.traincames
         if flag=='test':
            self.cames1= self.testcames
    def __len__(self):
        return len( self.cames1)*self.framnumber
    def __getitem__(self,index):
        cames=self.cames1[index//self.framnumber]
        fram=index%self.framnumber+1
        if fram==0:
            assert 0
        batch={}
        fram_ref=[-1,0,1]
        masks1=[]
        inputs1=[]
        dpoints=[]
        for i in fram_ref:
            mask,pictureL,pictureR,dpointL,dpointR,target=self.getdata(cames,fram+i)
            input1=torch.cat([pictureL,pictureR],axis=0)
            dpoint=torch.cat([dpointL,dpointR],axis=0)
            mask=1-mask.float()
            if i==0:
                targets=[target.float()]
            dpoints.append(dpoint.unsqueeze(0))
            masks1.append(mask.unsqueeze(0))
            inputs1.append(input1.unsqueeze(0))
        if(self.flag=='train'):
                targets,dpoints,masks1,inputs1= cut_picture(targets,dpoints,masks1,inputs1,dH=self.img_size,dW=self.img_size)
        batch['inputs']=torch.cat(inputs1,axis=0).float()                   #b,[ref],6,h,w
        batch['masks']=torch.cat(masks1,axis=0).float()
        batch['dpoints']=torch.cat(dpoints,axis=0).float()
        batch['target']=torch.cat(targets,axis=0).float()
        return batch